import cv2
import numpy as np
import time
from utils import capture_screenshot
from datetime import datetime

template_zai = cv2.imread("images/seal/zai.png", cv2.IMREAD_COLOR)
template_dao = cv2.imread("images/seal/dao.png", cv2.IMREAD_COLOR)

template_info_seal = [
    (template_zai, "zai", True),
    (template_dao, "dao", False),
]

template_pengzhang = cv2.imread("images/seal/template/pengzhang.png", cv2.IMREAD_COLOR)
template_shengming = cv2.imread("images/seal/template/shengming.png", cv2.IMREAD_COLOR)
template_xingfu = cv2.imread("images/seal/template/xingfu.png", cv2.IMREAD_COLOR)
template_dujin = cv2.imread("images/seal/template/dujin.png", cv2.IMREAD_COLOR)

template_info_seal_shengming = [
    (template_pengzhang, "pengzhang", True),
    (template_shengming, "shengming", True),
    (template_xingfu, "xingfu", True),
    (template_dujin, "dujin", True),
]

def find_seal():
    image = capture_screenshot(region_name="印章")
    numbers = find_numbers(image)
    seal_names = find_select_index_multi(image, template_info_seal_shengming, numbers)
    return seal_names

def find_numbers(image):
    """
    返回印章的颜色数组
    红色3 蓝色2 灰色1 其他0
    找不到则返回空
    :param image: 选项区的图片
    """
    # 固定的区域
    test_regions = [(0, 0, 1, 1), (468, 0, 1, 1), (937, 0, 1, 1)]

#     for (x, y, w, h) in test_regions:
#         # 在图片上绘制红色矩形框
#         cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255), 2)
#     save_path = "images/option/base2_t.png"
#     cv2.imwrite(save_path, image)

    average_colors = []
    for (x, y, w, h) in test_regions:
        # 提取指定区域
        roi = image[y:y + h, x:x + w]
        # 计算该区域的平均颜色
        average_color = np.mean(roi, axis=(0, 1))
        average_colors.append(average_color)

    gray_color = np.array([91, 101, 116])
    blue_color = np.array([186, 92, 109])
    red_color = np.array([15, 65, 240])
    error_margin = 30  # 误差余量

#     print(average_colors)
    numbers = []
    for color in average_colors:
        if np.all(np.abs(color - blue_color) <= error_margin):
            numbers.append(2)
        elif np.all(np.abs(color - red_color) <= error_margin):
            numbers.append(3)
        elif np.all(np.abs(color - gray_color) <= error_margin):
            numbers.append(1)
        else:
            numbers.append(0)
    return numbers

def find_select_index_multi(image, template_info, numbers):
    """
    选项截图中指定模板的位置
    返回匹配的位置（1、2、3）和模板名称
    找不到则返回 (0, None)
    :param image: 选项区的图片
    :param template_info: 模板信息列表，每个元素为 (模板, 模板名称, 开关条件)
    :param numbers: 用于控制保存图片的参数
    """
    # 固定的区域
    test_regions = [(0, 0, 51, 44), (468, 0, 51, 44), (937, 0, 51, 44)]
    # 提前提取指定区域
    rois = []
    for i, (x, y, w, h) in enumerate(test_regions):

        roi = image[y:y + h, x:x + w]
        rois.append(roi)
        # 根据 numbers 数组判断是否保存图片
        if numbers and 0 <= i < len(numbers) and numbers[i] > 0:
            now = datetime.now()
            current_time_str = now.strftime("%Y%m%d%H%M%S")
            save_path = f"images/seal/tem/{current_time_str}_{i + 1}.png"
            cv2.imwrite(save_path, roi)

    # 根据开关条件筛选模板
    valid_templates = [(name, template) for template, name, condition in template_info if condition]
    names = []
    for i, roi in enumerate(rois):
        name = None
        for template_name, template in valid_templates:
            result = cv2.matchTemplate(roi, template, cv2.TM_CCOEFF_NORMED)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
            if max_val > 0.7:
                name = template_name
        names.append(name)
    return names

if __name__ == "__main__":
#     test_save_path = "images/option/tem2/20250311194601_3.png"
#     image = cv2.imread(test_save_path)
#     test_save_path2 = "images/option/han_p.png"
#     image2 = cv2.imread(test_save_path2)
#     result = cv2.matchTemplate(image, image2, cv2.TM_CCOEFF_NORMED)
#     min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
#     print(f"相似度: {max_val}")
    name = find_seal()
    print(name)
#     create_template()

